import numpy as np
import util

from p01b_logreg import LogisticRegression

# Character to replace with sub-problem letter in plot_path/pred_path
WILDCARD = 'X'


def compute_probabilities(model, x):
    """计算逻辑回归模型的预测概率

    Args:
        model: 训练好的逻辑回归模型
        x: 输入特征，形状 (m, n)

    Returns:
        预测概率，形状 (m,)
    """
    # 计算线性组合
    z = np.dot(x, model.theta)
    # 应用sigmoid函数得到概率
    probabilities = 1 / (1 + np.exp(-z))
    return probabilities


def main(train_path, valid_path, test_path, pred_path):
    """Problem 2: Logistic regression for incomplete, positive-only labels.

    Run under the following conditions:
        1. on y-labels,
        2. on l-labels,
        3. on l-labels with correction factor alpha.

    Args:
        train_path: Path to CSV file containing training set.
        valid_path: Path to CSV file containing validation set.
        test_path: Path to CSV file containing test set.
        pred_path: Path to save predictions.
    """
    pred_path_c = pred_path.replace(WILDCARD, 'c')
    pred_path_d = pred_path.replace(WILDCARD, 'd')
    pred_path_e = pred_path.replace(WILDCARD, 'e')

    # *** START CODE HERE ***
    # Part (c): Train and test on true labels
    # 使用真实标签训练和测试
    print("Part (c): Training on true labels t")

    # 加载数据 - 使用util.py中的load_dataset函数
    # 根据题目要求：t是真实标签，y是标记标签
    x_train, t_train = util.load_dataset(train_path, label_col='t', add_intercept=True)
    x_valid, t_valid = util.load_dataset(valid_path, label_col='t', add_intercept=True)
    x_test, t_test = util.load_dataset(test_path, label_col='t', add_intercept=True)

    # 加载y标签（标记标签）
    _, y_train = util.load_dataset(train_path, label_col='y', add_intercept=False)
    _, y_valid = util.load_dataset(valid_path, label_col='y', add_intercept=False)
    _, y_test = util.load_dataset(test_path, label_col='y', add_intercept=False)

    # 训练逻辑回归模型（使用真实标签t）
    model_c = LogisticRegression()
    model_c.fit(x_train, t_train)

    # 在测试集上预测
    t_pred_c = model_c.predict(x_test)

    # 保存预测结果
    np.savetxt(pred_path_c, t_pred_c)

    # 计算准确率
    accuracy_c = np.mean(t_pred_c == t_test)
    print(f"Part (c) Accuracy: {accuracy_c:.4f}")

    # Part (d): Train on positive-only labels and test on true labels
    # 使用仅正例标签训练，在真实标签上测试
    print("\nPart (d): Training on positive-only labels...")

    # 使用标记标签 y 训练模型
    # y = 1 表示被标记的样本，y = 0 表示未被标记的样本
    model_d = LogisticRegression()
    model_d.fit(x_train, y_train)

    # 在测试集上预测
    t_pred_d = model_d.predict(x_test)

    # 保存预测结果
    np.savetxt(pred_path_d, t_pred_d)

    # 计算准确率
    accuracy_d = np.mean(t_pred_d == t_test)
    print(f"Part (d) Accuracy: {accuracy_d:.4f}")

    # Part (e): Apply correction factor using validation set and test on true labels
    # 使用验证集应用修正因子
    print("\nPart (e): Applying correction factor：")

    # 找到验证集中被标记的样本
    V_plus_indices = y_valid == 1  # 被标记的样本
    V_plus = x_valid[V_plus_indices]  # 被标记样本的特征

    print(f"验证集被标记样本数量: {np.sum(V_plus_indices)}")

    if np.sum(V_plus_indices) > 0:
        # 使用(d)部分的分类器计算在被标记样本上的预测概率
        h_V_plus = compute_probabilities(model_d, V_plus)

        alpha = np.mean(h_V_plus)

        print(f"α  = {alpha:.4f}")

        # 使用(d)部分的分类器对测试集进行预测，然后用α修正
        test_probs = compute_probabilities(model_d, x_test)
        corrected_probs = test_probs / alpha
        corrected_probs = np.clip(corrected_probs, 0, 1)

        # 使用阈值0.5进行预测：p(t=1|x) = 0.5
        t_pred_e = (corrected_probs > 0.5).astype(int)

    else:
        print(f"警告: 验证集中没有被标记的样本，无法计算α")
        t_pred_e = t_pred_d
        alpha = 1.0

    # 保存预测结果
    np.savetxt(pred_path_e, t_pred_e)

    # 计算准确率
    accuracy_e = np.mean(t_pred_e == t_test)
    print(f"Part (e) Accuracy: {accuracy_e:.4f}")

    # 绘制三种方法的决策边界
    util.plot(x_test, t_test, model_c.theta, pred_path_c.replace('.txt', '.png'))
    util.plot(x_test, t_test, model_d.theta, pred_path_d.replace('.txt', '.png'))

    # 方法(e)的决策边界：使用修正后的theta
    if alpha != 1.0:
        # 计算修正后的决策边界阈值
        corrected_threshold = np.log(0.5 * alpha / (1 - 0.5 * alpha))
        corrected_theta = model_d.theta.copy()
        corrected_theta[0] = corrected_theta[0] - corrected_threshold
        util.plot(x_test, t_test, corrected_theta, pred_path_e.replace('.txt', '.png'))
    else:
        util.plot(x_test, t_test, model_d.theta, pred_path_e.replace('.txt', '.png'))

    # *** END CODE HERE ***